{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/peremartra/Large-Language-Model-Notebooks-Course/blob/main/6-PRUNING/6_1_pruning_structured_l1_diltilgpt2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uMp3lltu2F_N"
      },
      "source": [
        "\n",
        "<div>\n",
        "    <h1>Large Language Models Projects</a></h1>\n",
        "    <h3>Apply and Implement Strategies for Large Language Models</h3>\n",
        "    <h2>Pruning distilGPT2.</h2>\n",
        "    <h3>Structured Width Pruning: Eliminating Less Important Neurons from Feedforward Layers.</h3>\n",
        "</div>\n",
        "\n",
        "by [Pere Martra](https://www.linkedin.com/in/pere-martra/)\n",
        "_______\n",
        "Models: distilgpt2\n",
        "\n",
        "Colab Environment: CPU / GPU T4.\n",
        "\n",
        "Keys:\n",
        "* Pruning\n",
        "* Structured pruning\n",
        "\n",
        "\n",
        "Related article: --.\n",
        "_______\n",
        "**disclaimer: The pruning section was created after the first edition of the book was published. They are not included in the book’s original content but are intended to supplement and expand on the topics covered.**\n",
        "\n",
        "\n",
        "This is the unofficial repository for the book:\n",
        "        <a href=\"https://amzn.to/4eanT1g\"> <b>Large Language Models:</b> Apply and Implement Strategies for Large Language Models</a> (Apress).\n",
        "        The book is based on the content of this repository, but the notebooks are being updated, and I am incorporating new examples and chapters.\n",
        "        If you are looking for the official repository for the book, with the original notebooks, you should visit the\n",
        "        <a href=\"https://github.com/Apress/Large-Language-Models-Projects\">Apress repository</a>, where you can find all the notebooks in their original format as they appear in the book."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pd_dMCGn3bC4"
      },
      "source": [
        "#PRUNING\n",
        "## neurons structured width pruning\n",
        "Pruning is an important optimization technique in machine learning that aims to reduce the size of a model without sacrificing much of its accuracy. By removing less important components, pruning not only decreases the computational cost but also makes the model more efficient for deployment, especially on resource-constrained devices.\n",
        "\n",
        "Can be compared to quantization, another optimization technique that reduces the precision of the model's weights, typically converting them from high-precision floating-point numbers to lower-precision representations. While quantization can significantly reduce model size and speed up inference, it does not selectively remove weights.\n",
        "\n",
        "On the other hand, pruning, allows for targeted removal of less important weights or neurons, which can lead to a more efficient reduction in model size while better preserving accuracy. By selecting the weights to eliminate based on their importance scores, pruning provides more control over the model's structure, often making it a more effective approach when aiming for both model compression and high performance.\n",
        "\n",
        "The effectiveness of removing specific parts of a model could be debated, but recent studies, such as the one conducted by NVIDIA: [How to Prune and Distill Llama-3.1 8B to an NVIDIA Llama-3.1-Minitron 4B Model](https://developer.nvidia.com/blog/how-to-prune-and-distill-llama-3-1-8b-to-an-nvidia-llama-3-1-minitron-4b-model/), concluding that pruning, combined with fine-tuning techniques applied after pruning, can produce models that are not only more efficient but also more effective in specific domains.\n",
        "\n",
        "Also you can combine both techniques and quantize a model that has been previousluy pruned.\n",
        "\n",
        "This notebook focuses on **structured width pruning**, where entire neurons are eliminated based on their low importance scores, which are computed using the L1 norm. The assumption is that neurons with lower L1 norm values contribute less to the overall output of the model, allowing for safe removal to enhance efficiency without drastically impacting accuracy."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pwXWapcuhMVV"
      },
      "source": [
        "# Install Libraries & Configure variables."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "JfE4lAap1VSi",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0c22ac66-e03a-40d5-b81f-aa916af11f8e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m363.4/363.4 MB\u001b[0m \u001b[31m4.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.8/13.8 MB\u001b[0m \u001b[31m92.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m24.6/24.6 MB\u001b[0m \u001b[31m80.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m883.7/883.7 kB\u001b[0m \u001b[31m45.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m664.8/664.8 MB\u001b[0m \u001b[31m2.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m211.5/211.5 MB\u001b[0m \u001b[31m6.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.3/56.3 MB\u001b[0m \u001b[31m11.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m127.9/127.9 MB\u001b[0m \u001b[31m7.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m207.5/207.5 MB\u001b[0m \u001b[31m6.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.1/21.1 MB\u001b[0m \u001b[31m91.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "#Install necessary libraries\n",
        "!pip install -q transformers\n",
        "!pip install -q torch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f6svKMvV1fG8",
        "outputId": "f7fd1ccf-b452-4a22-c06e-2c9ffd9ba9a1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using device: cpu\n"
          ]
        }
      ],
      "source": [
        "#Import libraries\n",
        "import torch\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "from torch import nn\n",
        "import os\n",
        "\n",
        "# Check if GPU is available\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print(f\"Using device: {device}\")\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FK88hwK1C0NJ"
      },
      "source": [
        "I chose to prune only 20% of the least important neurons based on their L1 norm, aiming to balance size reduction with minimal accuracy loss.\n",
        "\n",
        "You can adjust this percentage and increase it to 30% or even 50%, depending on whether you plan to follow up with a post-pruning fine-tuning process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "Duz5JbGYd0Jm"
      },
      "outputs": [],
      "source": [
        "prune_percent = 0.2  # Prune 20% of neurons\n",
        "model_name = 'distilgpt2'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "v0fIEu2jeH10"
      },
      "outputs": [],
      "source": [
        "# Support function to check the size reduction.\n",
        "def count_parameters(model):\n",
        "    return sum(p.numel() for p in model.parameters())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SU87XHJHX18o"
      },
      "source": [
        "#Download Model and explore structure."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oyWPcQoy1nKi"
      },
      "outputs": [],
      "source": [
        "# Download the model and tokenizer\n",
        "model = AutoModelForCausalLM.from_pretrained(model_name).to(device)\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "tokenizer.pad_token = tokenizer.eos_token\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "fq-WAX4MfHbd"
      },
      "outputs": [],
      "source": [
        "def get_ouput(prompt, model=model, tokenizer=tokenizer):\n",
        "  inputs = tokenizer(prompt, return_tensors='pt').to(device)\n",
        "  outputs = model.generate(inputs['input_ids'],\n",
        "                           attention_mask=inputs['attention_mask'],\n",
        "                           max_length=10,\n",
        "                           num_return_sequences=1)\n",
        "  generated = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
        "  return generated"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RPguo_kIKb2g"
      },
      "source": [
        "## studying the model's structure.\n",
        "\n",
        "Understanding the model's structure is crucial in a pruning process.\n",
        "\n",
        "In this structure, you can see the part dedicated to the Attention layers (attn) and the part dedicated to the FeedForward layers (mlp).\n",
        "\n",
        "In the pruning process for the notebook, I only targeted the mlp layers because they typically contribute the most to the model's size and pruning them doesn’t affect the attention mechanism, which is critical for capturing relationships within the input data. These layers also tend to contain more redundancies, and reducing neurons in them generally doesn't significantly impact the model's output—though this always depends on the neuron selection process for elimination.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B4NhrhqfX749",
        "outputId": "7c3b85f9-4c77-4815-f281-6fa6f4f1748a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GPT2LMHeadModel(\n",
            "  (transformer): GPT2Model(\n",
            "    (wte): Embedding(50257, 768)\n",
            "    (wpe): Embedding(1024, 768)\n",
            "    (drop): Dropout(p=0.1, inplace=False)\n",
            "    (h): ModuleList(\n",
            "      (0-5): 6 x GPT2Block(\n",
            "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPT2Attention(\n",
            "          (c_attn): Conv1D(nf=2304, nx=768)\n",
            "          (c_proj): Conv1D(nf=768, nx=768)\n",
            "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
            "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
            "        )\n",
            "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPT2MLP(\n",
            "          (c_fc): Conv1D(nf=3072, nx=768)\n",
            "          (c_proj): Conv1D(nf=768, nx=3072)\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.1, inplace=False)\n",
            "        )\n",
            "      )\n",
            "    )\n",
            "    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "  )\n",
            "  (lm_head): Linear(in_features=768, out_features=50257, bias=False)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "print(model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YZSjvokoQtCT"
      },
      "source": [
        "In the model's structure, we find the embedding layers: `wte` (Word Token Embedding) and `wpe` (Word Position Embedding). The vector used to represent the input data has a size of 768.\n",
        "\n",
        "After the embedding layers, there's a dropout layer.\n",
        "\n",
        "Next, we have the typical layers of a Transformer model:\n",
        "\n",
        "- **Normalization layers** (`ln_1` and `ln_2`).\n",
        "- **Attention mechanism** (`attn`), consisting of its convolutional layers and dropout layers.\n",
        "- **Feed-forward layers** (`mlp), which will be the target of the pruning process. Specifically, I have chosen to prune the `c_fc` and `c_proj` layers. These layers expand and compress the information that passes through them. They are necessary for enabling the model to capture complex relationships within the input data, but it's quite easy to find neurons in these layers that don't contribute much to the model, at least when using the model with specific data.\n",
        "- The model ends with the **final normalization layer** (`ln_f`)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z2hXvLVnSScu"
      },
      "source": [
        "Another important consideration is the model's configuration file. Since the pruning process alters the model's structure, the resulting structure must be reflected in the configuration file.\n",
        "\n",
        "Otherwise, we might encounter issues where the model doesn't work properly with the Transformers library or produces errors or incorrect results during inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6sVAQPPGX_RQ",
        "outputId": "0e4dd1b6-6fc5-49cf-e66a-3bedaf617bde"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GPT2Config {\n",
            "  \"_attn_implementation_autoset\": true,\n",
            "  \"_num_labels\": 1,\n",
            "  \"activation_function\": \"gelu_new\",\n",
            "  \"architectures\": [\n",
            "    \"GPT2LMHeadModel\"\n",
            "  ],\n",
            "  \"attn_pdrop\": 0.1,\n",
            "  \"bos_token_id\": 50256,\n",
            "  \"embd_pdrop\": 0.1,\n",
            "  \"eos_token_id\": 50256,\n",
            "  \"id2label\": {\n",
            "    \"0\": \"LABEL_0\"\n",
            "  },\n",
            "  \"initializer_range\": 0.02,\n",
            "  \"label2id\": {\n",
            "    \"LABEL_0\": 0\n",
            "  },\n",
            "  \"layer_norm_epsilon\": 1e-05,\n",
            "  \"model_type\": \"gpt2\",\n",
            "  \"n_ctx\": 1024,\n",
            "  \"n_embd\": 768,\n",
            "  \"n_head\": 12,\n",
            "  \"n_inner\": null,\n",
            "  \"n_layer\": 6,\n",
            "  \"n_positions\": 1024,\n",
            "  \"reorder_and_upcast_attn\": false,\n",
            "  \"resid_pdrop\": 0.1,\n",
            "  \"scale_attn_by_inverse_layer_idx\": false,\n",
            "  \"scale_attn_weights\": true,\n",
            "  \"summary_activation\": null,\n",
            "  \"summary_first_dropout\": 0.1,\n",
            "  \"summary_proj_to_labels\": true,\n",
            "  \"summary_type\": \"cls_index\",\n",
            "  \"summary_use_proj\": true,\n",
            "  \"task_specific_params\": {\n",
            "    \"text-generation\": {\n",
            "      \"do_sample\": true,\n",
            "      \"max_length\": 50\n",
            "    }\n",
            "  },\n",
            "  \"torch_dtype\": \"float32\",\n",
            "  \"transformers_version\": \"4.51.3\",\n",
            "  \"use_cache\": true,\n",
            "  \"vocab_size\": 50257\n",
            "}\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(model.config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VthKlKdZSPkd",
        "outputId": "b47829ef-93ee-45f9-ec38-84483147c6d9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated text: Paris is the capital of the United States.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "#Test the original model with a simple prompt\n",
        "prompt = \"Paris is the capital of\"\n",
        "generated = get_ouput(prompt)\n",
        "print(f\"Generated text: {generated}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pGjseuY42h9O",
        "outputId": "ec023346-0a53-4d65-b9a0-39ac8c197eac"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original model parameters: 81912576\n"
          ]
        }
      ],
      "source": [
        "#Print the size of the original model\n",
        "original_param_count = count_parameters(model)\n",
        "print(f\"Original model parameters: {original_param_count}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mB2stu4J1azJ"
      },
      "source": [
        "# Pruning Model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g1MIzdVriHlj"
      },
      "source": [
        "The MLP layers in GPT-2 use 1D convolutions (Conv1D) for their transformations. This line imports the Conv1D class from the GPT-2 model implementation, which will be used later to create new layers with reduced sizes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "5oZ2K2z6h9Fi"
      },
      "outputs": [],
      "source": [
        "#Prune the MLP layers based on weight magnitude (adjusted for Conv1D layers)\n",
        "from transformers.models.gpt2.modeling_gpt2 import Conv1D\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TYMe1zBTiQzB"
      },
      "source": [
        "This variable is used to store the new intermediate size of the MLP after pruning. It will be updated once the number of neurons to keep is determined."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "0_LCgVAWiTHa"
      },
      "outputs": [],
      "source": [
        "# Initialize new_intermediate_size\n",
        "new_intermediate_size = None"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "byVBlNjOi9yL"
      },
      "source": [
        "## Support pruing functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "Ert9JitClPAu"
      },
      "outputs": [],
      "source": [
        "# Function to compute importance scores (L1 norm)\n",
        "def compute_importance_scores(c_fc_weight):\n",
        "    \"\"\"\n",
        "    Compute the importance scores for each neuron in the c_fc layer using L1 norm.\n",
        "\n",
        "    Args:\n",
        "    - c_fc_weight: Weight matrix from the c_fc layer.\n",
        "\n",
        "    Returns:\n",
        "    - importance_scores: L1 norm importance scores for each neuron.\n",
        "    \"\"\"\n",
        "    return torch.sum(torch.abs(c_fc_weight), dim=0)  # Shape: [intermediate_size]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "_-tFAMWslQ8m"
      },
      "outputs": [],
      "source": [
        "# Function to prune neurons and create new Conv1D layers\n",
        "def prune_neurons(mlp, prune_percent, device):\n",
        "    \"\"\"\n",
        "    Prune neurons from the c_fc and c_proj layers of the MLP based on importance scores.\n",
        "\n",
        "    Args:\n",
        "    - mlp: The MLP layer (contains c_fc and c_proj) to prune.\n",
        "    - prune_percent: Percentage of neurons to prune.\n",
        "    - device: Device (CPU/GPU) for model operations.\n",
        "\n",
        "    Returns:\n",
        "    - new_c_fc: New pruned c_fc layer.\n",
        "    - new_c_proj: New pruned c_proj layer.\n",
        "    - new_intermediate_size: Size of the pruned intermediate layer.\n",
        "    \"\"\"\n",
        "    # Get the weights of the c_fc layer (input projection)\n",
        "    c_fc_weight = mlp.c_fc.weight.data\n",
        "\n",
        "    # Compute importance scores for each neuron\n",
        "    importance_scores = compute_importance_scores(c_fc_weight)\n",
        "\n",
        "    # Determine the number of neurons to prune\n",
        "    original_intermediate_size = c_fc_weight.size(1)  # This is intermediate_size\n",
        "    num_neurons_to_prune = int(prune_percent * original_intermediate_size)\n",
        "\n",
        "    # Get indices of neurons to keep (those with highest importance)\n",
        "    _, indices_to_keep = torch.topk(importance_scores, original_intermediate_size - num_neurons_to_prune)\n",
        "\n",
        "    # Sort indices to maintain order\n",
        "    indices_to_keep, _ = torch.sort(indices_to_keep)\n",
        "\n",
        "    # Create new Conv1D layers with reduced size\n",
        "    new_c_fc = Conv1D(len(indices_to_keep), mlp.c_fc.weight.size(0)).to(device)  # Conv1D(new_intermediate_size, hidden_size)\n",
        "    new_c_proj = Conv1D(mlp.c_proj.weight.size(1), len(indices_to_keep)).to(device)  # Conv1D(hidden_size, new_intermediate_size)\n",
        "\n",
        "    return new_c_fc, new_c_proj, len(indices_to_keep), indices_to_keep"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "EIG9ld8JlVYk"
      },
      "outputs": [],
      "source": [
        "# Function to copy weights and biases to new pruned layers\n",
        "def copy_weights_and_biases(mlp, new_c_fc, new_c_proj, indices_to_keep):\n",
        "    \"\"\"\n",
        "    Copy the weights and biases from the original layers to the new pruned layers.\n",
        "\n",
        "    Args:\n",
        "    - mlp: The original MLP layer (contains c_fc and c_proj).\n",
        "    - new_c_fc: New pruned c_fc layer.\n",
        "    - new_c_proj: New pruned c_proj layer.\n",
        "    - indices_to_keep: Indices of neurons that are retained.\n",
        "    \"\"\"\n",
        "    # Copy weights and biases for the neurons we are keeping\n",
        "    new_c_fc.weight.data = mlp.c_fc.weight.data[:, indices_to_keep]\n",
        "    new_c_fc.bias.data = mlp.c_fc.bias.data[indices_to_keep]\n",
        "\n",
        "    new_c_proj.weight.data = mlp.c_proj.weight.data[indices_to_keep, :]\n",
        "    new_c_proj.bias.data = mlp.c_proj.bias.data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EQU00eFSY7QU"
      },
      "source": [
        "## Prune Loop"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KAadbLUBLB5m"
      },
      "source": [
        "\n",
        "The `update_model` function iterates through the blocks within the model's Transformer structure. This structure consists of multiple `GPT2Block` blocks, and each of these blocks contains a pair of `GPT2SdpaAttention` and `GPT2MLP` components. The latter contains the MLP layers that will be the target of the pruning process.\n",
        "```\n",
        "(h): ModuleList(\n",
        "      (0-5): 6 x GPT2Block(\n",
        "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
        "        (attn): GPT2SdpaAttention(\n",
        "          (c_attn): Conv1D()\n",
        "          (c_proj): Conv1D()\n",
        "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
        "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
        "        )\n",
        "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
        "        (mlp): GPT2MLP(\n",
        "          (c_fc): Conv1D()\n",
        "          (c_proj): Conv1D()\n",
        "          (act): NewGELUActivation()\n",
        "          (dropout): Dropout(p=0.1, inplace=False)\n",
        "        )\n",
        "      )\n",
        "    )\n",
        "```\n",
        "\n",
        "The layers that will undergo the removal of neurons identified as less useful are:\n",
        "\n",
        "**(c_fc): Conv1D()**\n",
        "\n",
        "**(c_proj): Conv1D()**\n",
        "\n",
        "The neurons are removed in the `prune_neurons` function based on the values returned by `compute_importance_scores`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "-v5na0YnlF0d"
      },
      "outputs": [],
      "source": [
        "# Function to iterate through the model and prune each block\n",
        "def update_model(model, prune_percent, device):\n",
        "    \"\"\"\n",
        "    Prune the MLP layers of each Transformer block in the model and update the model's configuration.\n",
        "\n",
        "    Args:\n",
        "    - model: The GPT-2 model to prune.\n",
        "    - prune_percent: Percentage of neurons to prune.\n",
        "    - device: Device (CPU/GPU) for model operations.\n",
        "\n",
        "    Returns:\n",
        "    - model: The pruned model with updated layers.\n",
        "    - new_intermediate_size: The new intermediate size after pruning.\n",
        "    \"\"\"\n",
        "    new_intermediate_size = None\n",
        "\n",
        "    # Iterate through each block in the model\n",
        "    for idx, block in enumerate(model.transformer.h):\n",
        "        mlp = block.mlp\n",
        "\n",
        "        # Prune the neurons and create new layers\n",
        "        new_c_fc, new_c_proj, new_size, indices_to_keep = prune_neurons(mlp, prune_percent, device)\n",
        "\n",
        "        # Copy weights and biases from old layers to new pruned layers\n",
        "        copy_weights_and_biases(mlp, new_c_fc, new_c_proj, indices_to_keep)\n",
        "\n",
        "        # Replace old layers with new pruned layers\n",
        "        mlp.c_fc = new_c_fc\n",
        "        mlp.c_proj = new_c_proj\n",
        "\n",
        "        # Update the intermediate size for the first block\n",
        "        if new_intermediate_size is None:\n",
        "            new_intermediate_size = new_size\n",
        "\n",
        "    # Update the model configuration with the new intermediate size\n",
        "    model.config.n_inner = new_intermediate_size\n",
        "\n",
        "    return model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LqtGqDTy1lbt"
      },
      "source": [
        "## Obtain & Check pruned model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "y6PSMIe0lYt8"
      },
      "outputs": [],
      "source": [
        "# Get the pruned Model\n",
        "model = update_model(model, prune_percent, device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "51hobl7W4iou",
        "outputId": "50d285e1-2e8c-4de6-da51-6c0f631d7632"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Pruned model parameters: 76250268\n",
            "Reduction in parameters: 5662308\n"
          ]
        }
      ],
      "source": [
        "#Step 7: Recalculate the number of parameters\n",
        "pruned_param_count = count_parameters(model)\n",
        "print(f\"Pruned model parameters: {pruned_param_count}\")\n",
        "print(f\"Reduction in parameters: {original_param_count - pruned_param_count}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u2sGEFIbZbnr"
      },
      "source": [
        "The savings produced by the pruning process is around 7.5%. It might seem like a small reward for all the effort, but we can adjust the percentage of pruned neurons. More importantly, we can achieve a more efficient model than the base through a subsequent fine-tuning process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E_oRTPuD1p-D",
        "outputId": "d8528308-9647-45bc-c0a5-e7cf7caa4a35"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GPT2LMHeadModel(\n",
            "  (transformer): GPT2Model(\n",
            "    (wte): Embedding(50257, 768)\n",
            "    (wpe): Embedding(1024, 768)\n",
            "    (drop): Dropout(p=0.1, inplace=False)\n",
            "    (h): ModuleList(\n",
            "      (0-5): 6 x GPT2Block(\n",
            "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPT2Attention(\n",
            "          (c_attn): Conv1D(nf=2304, nx=768)\n",
            "          (c_proj): Conv1D(nf=768, nx=768)\n",
            "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
            "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
            "        )\n",
            "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPT2MLP(\n",
            "          (c_fc): Conv1D(nf=2458, nx=768)\n",
            "          (c_proj): Conv1D(nf=768, nx=2458)\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.1, inplace=False)\n",
            "        )\n",
            "      )\n",
            "    )\n",
            "    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "  )\n",
            "  (lm_head): Linear(in_features=768, out_features=50257, bias=False)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "#structure prompted model.\n",
        "print(model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AxnDB9BramIF"
      },
      "source": [
        "The model's structure after going through the pruning process appears unchanged. This is because I have only removed neurons, not entire layers, so the layer weights have been altered, but the layers themselves remain intact.\n",
        "\n",
        "Torch shows the layers but not their internal weights, which is why the structure seems the same. However, when counting the model's parameters, you can see a reduction of 5,662,308 parameters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "89AbYZ5F1vxA",
        "outputId": "da9aae69-2624-41b4-c29f-0cffbad86f0d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GPT2Config {\n",
            "  \"_attn_implementation_autoset\": true,\n",
            "  \"_num_labels\": 1,\n",
            "  \"activation_function\": \"gelu_new\",\n",
            "  \"architectures\": [\n",
            "    \"GPT2LMHeadModel\"\n",
            "  ],\n",
            "  \"attn_pdrop\": 0.1,\n",
            "  \"bos_token_id\": 50256,\n",
            "  \"embd_pdrop\": 0.1,\n",
            "  \"eos_token_id\": 50256,\n",
            "  \"id2label\": {\n",
            "    \"0\": \"LABEL_0\"\n",
            "  },\n",
            "  \"initializer_range\": 0.02,\n",
            "  \"label2id\": {\n",
            "    \"LABEL_0\": 0\n",
            "  },\n",
            "  \"layer_norm_epsilon\": 1e-05,\n",
            "  \"model_type\": \"gpt2\",\n",
            "  \"n_ctx\": 1024,\n",
            "  \"n_embd\": 768,\n",
            "  \"n_head\": 12,\n",
            "  \"n_inner\": 2458,\n",
            "  \"n_layer\": 6,\n",
            "  \"n_positions\": 1024,\n",
            "  \"reorder_and_upcast_attn\": false,\n",
            "  \"resid_pdrop\": 0.1,\n",
            "  \"scale_attn_by_inverse_layer_idx\": false,\n",
            "  \"scale_attn_weights\": true,\n",
            "  \"summary_activation\": null,\n",
            "  \"summary_first_dropout\": 0.1,\n",
            "  \"summary_proj_to_labels\": true,\n",
            "  \"summary_type\": \"cls_index\",\n",
            "  \"summary_use_proj\": true,\n",
            "  \"task_specific_params\": {\n",
            "    \"text-generation\": {\n",
            "      \"do_sample\": true,\n",
            "      \"max_length\": 50\n",
            "    }\n",
            "  },\n",
            "  \"torch_dtype\": \"float32\",\n",
            "  \"transformers_version\": \"4.51.3\",\n",
            "  \"use_cache\": true,\n",
            "  \"vocab_size\": 50257\n",
            "}\n",
            "\n"
          ]
        }
      ],
      "source": [
        "#config file pruned model.\n",
        "print(model.config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C82nbgh1coTV"
      },
      "source": [
        "In the configuration file, a difference is noticeable: the n_inner parameter now contains the value representing the number of neurons in the c_fc layer, a feedforward layer that has had its neuron count reduced.\n",
        "\n",
        "Not all models handle the information from this file the same way, but in the case of the distilgpt2 model, if the n_inner value is null, the default value is set to four times the hidden_size. We can see the size in the model's structure within the embedding layers.\n",
        "\n",
        "In this case, the default n_inner value would be 4 * 768 = 3072, but since the weights of the layers have been reduced through the pruning process, it has been replaced with 2458."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "S2gAXLSR_arg",
        "outputId": "23ae8695-4fda-4bc4-c26e-3c02affb56fe"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated text after pruning: Paris is the capital of the United States, and\n"
          ]
        }
      ],
      "source": [
        "generated = get_ouput(prompt)\n",
        "print(f\"Generated text after pruning: {generated}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nq7p1NbAey0c"
      },
      "source": [
        "The response from the pruned model differs from that of the base model, indicating that the pruning process has affected the model's output generation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gZjDxKrVbVo6"
      },
      "source": [
        "# Upload the model to Hugging Face & Download to test.\n",
        "\n",
        "We cannot be sure that the model works correctly with the Transformers library until we complete a full test cycle with it.\n",
        "\n",
        "Often, if the configuration file has not been properly modified, the issue arises during the process of downloading the model file from Hugging Face."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eYG4Zd55_hXg",
        "outputId": "09e68b70-0783-4555-c8bc-7d8e21ea8f40"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Pruned model saved to ./pruned_distilgpt2\n"
          ]
        }
      ],
      "source": [
        "# Save the pruned model\n",
        "output_dir = './pruned_distilgpt2'\n",
        "if not os.path.exists(output_dir):\n",
        "    os.makedirs(output_dir)\n",
        "\n",
        "model.save_pretrained(output_dir)\n",
        "tokenizer.save_pretrained(output_dir)\n",
        "print(f\"Pruned model saved to {output_dir}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zGGk--ndA1LA"
      },
      "outputs": [],
      "source": [
        "# Push the model to your Hugging Face repository\n",
        "name_model_to_push=\"pruned_distilgpt2\"\n",
        "\n",
        "model.push_to_hub(name_model_to_push,\n",
        "                  private=True,\n",
        "                  use_temp_dir=False)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n-LSLfjFBDHt"
      },
      "outputs": [],
      "source": [
        "tokenizer.push_to_hub(name_model_to_push,\n",
        "                      private=False,\n",
        "                      use_temp_dir=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BBWXDT0FMF2e"
      },
      "outputs": [],
      "source": [
        "# Step 11: Download the model from Hugging Face\n",
        "pruned_model_name = 'oopere/pruned_distilgpt2'\n",
        "pruned_model = AutoModelForCausalLM.from_pretrained(pruned_model_name).to(device)\n",
        "pruned_tokenizer = AutoTokenizer.from_pretrained(pruned_model_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lwmyUxXVNLBc",
        "outputId": "dba88a97-eec5-4141-c1cc-2a80b754c1dd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Pruned Downloaded Generated text: Paris is the capital of the United States, and\n"
          ]
        }
      ],
      "source": [
        "generated = get_ouput(prompt, pruned_model, pruned_tokenizer)\n",
        "print(f\"Pruned Downloaded Generated text: {generated}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9W_DVjFDb_Eg"
      },
      "source": [
        "# Conclusion.\n",
        "\n",
        "In this notebook, a complete pruning process has been applied to a Transformer model.\n",
        "\n",
        "The pruning process followed is a structured approach that removes specific neurons from the feedforward layers of the model.\n",
        "\n",
        "This results in a smaller model compared to the original, while retaining much of its ability to understand data relationships, as the attention layers remain untouched, and its learning capabilities remain intact since no layers were removed.\n",
        "\n",
        "The neurons eliminated are those deemed less important for the model's output. This method of selecting neurons, without using a dataset, is ideal when the goal is to obtain a model capable of mimicking the base model’s responses.\n",
        "\n",
        "It was taken into account that the pruning process modified the size of the layers, so the model’s configuration file had to be adjusted accordingly to ensure it continues to function without issues.\n",
        "\n",
        "## Future Work."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vXpqooySklO8"
      },
      "source": [
        "There are many different ways to continue building upon the work done. The two main approaches could be:\n",
        "\n",
        "* Use a dataset to select the least important neurons. This method allows the model to be adapted to a specific task, reducing its size and potentially increasing efficiency without requiring a subsequent fine-tuning process.\n",
        "* Perform depth pruning by removing entire layers from the model, rather than just specific neurons.\n",
        "* Adapt the approach to a larger and more current model, such as those from Meta's LLaMA family, Google's Gemma, or any other.\n",
        "\n",
        "It’s important to keep in mind that most models undergoing a pruning process are often fine-tuned afterward to regain the effectiveness they might have lost during pruning.\n",
        "\n",
        "The type of pruning to apply often depends on the task for which the model is being trained. For instance, one could consider pruning the attention layers if the input prompts are very short and the relationship between the tokens in the prompt is not particularly important, perhaps in tasks like classification or entity recognition.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F8bNxCJBqfNm"
      },
      "source": [
        "## Author's Note\n",
        "\n",
        "In addition to creating content like this notebook and offering it under the MIT license, I have also contributed to repositories such as those of Hugging Face and Google Gemini.\n",
        "\n",
        "I am especially proud of my book: <a href=\"https://amzn.to/4eanT1g\"><b>Large Language Models:</b> Apply and Implement Strategies for Large Language Models</a> (Apress).\n",
        "\n",
        "You can find it on both <a href=\"https://amzn.to/4eanT1g\">Amazon</a> and <a href=\"https://link.springer.com/book/10.1007/979-8-8688-0515-8\">Springer</a>, where they often have good deals on the purchase price.\n",
        "\n",
        "If you take a look and end up purchasing it, keep in mind that you can reach out with any questions via the Discussions section of this same repository or on any of my social media channels. I’ll do my best to respond as quickly as possible.\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPd+iItCMuNXwFdbYTxTePe",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}